﻿
#include "NextH264EncoderCPU.h"
#include "ffmpegdl.h"

namespace videonext
{
	class H264EncoderCPUVars : public Vars
	{
	public:
		Pointer encoder;
		Pointer frame;
		Pointer packet;
		Pointer bgr2yuv420p;
		UInt sendCount;
		Size2D size;
		Stream<Double> times;
		Stream<Binary> packets;

		~H264EncoderCPUVars()
		{
			close();
		}

		void close()
		{
			if (encoder)
			{
				avcodec_free_context(&encoder);
				encoder = 0;
			}
			if (frame)
			{
				av_frame_free(&frame);
				frame = 0;
			}
			if (packet)
			{
				if (getFFMPEGLibraryVersion() >= 3)
				{
					av_packet_free(&packet);
				}
				else
				{
					av_free_packet(packet);
				}
				packet = 0;
			}
			if (bgr2yuv420p)
			{
				sws_freeContext(bgr2yuv420p);
				bgr2yuv420p = 0;
			}
		}
	};
}

using namespace videonext;

H264EncoderCPU::H264EncoderCPU() : Object<H264EncoderCPUVars>(new H264EncoderCPUVars, TRUE)
{
	vars->encoder = 0;
	vars->frame = 0;
	vars->packet = 0;
	vars->bgr2yuv420p = 0;
}
Bool H264EncoderCPU::isSupported()
{
	Pointer targetCodec = avcodec_find_encoder_by_name("libx264");
	return targetCodec != 0;
}
Bool H264EncoderCPU::open(Size2D frameSize)
{
	vars->close();

	Pointer targetCodec = avcodec_find_encoder_by_name("libx264");
	if (!targetCodec) return FALSE;

	vars->encoder = avcodec_alloc_context3(targetCodec);
	if (!vars->encoder) return FALSE;

	avcodec_set_bit_rate(vars->encoder, 0);
	avcodec_set_width(vars->encoder, frameSize.width);
	avcodec_set_height(vars->encoder, frameSize.height);
	avcodec_set_time_base(vars->encoder, 1, 30);
	avcodec_set_framerate(vars->encoder, 30, 1);
	avcodec_set_gop_size(vars->encoder, 32);
	avcodec_set_max_b_frames(vars->encoder, 1);
	avcodec_set_pix_fmt(vars->encoder, AVPixelFormat::AV_PIX_FMT_YUV420P);

	av_opt_set(avcodec_get_priv_data(vars->encoder), "preset", "ultrafast", 0);
	av_opt_set(avcodec_get_priv_data(vars->encoder), "tune", "zerolatency,fastdecode", 0);

	Int resultOpenCodec = avcodec_open2(vars->encoder, targetCodec, NULL);
	if (resultOpenCodec < 0) return FALSE;

	vars->frame = av_frame_alloc();
	if (!vars->frame) return FALSE;

	av_frame_set_format(vars->frame, avcodec_get_pix_fmt(vars->encoder));
	av_frame_set_width(vars->frame, avcodec_get_width(vars->encoder));
	av_frame_set_height(vars->frame, avcodec_get_height(vars->encoder));

	Int resultCreateFrameData = av_frame_get_buffer(vars->frame, getFFMPEGLibraryVersion() >= 3 ? 0 : 8);
	if (resultCreateFrameData < 0) return FALSE;

	vars->bgr2yuv420p = sws_getContext(frameSize.width, frameSize.height, AVPixelFormat::AV_PIX_FMT_BGR24, frameSize.width, frameSize.height, avcodec_get_pix_fmt(vars->encoder), SWS_POINT, NULL, NULL, NULL);
	if (!vars->bgr2yuv420p) return FALSE;

	if (getFFMPEGLibraryVersion() >= 3)
	{
		vars->packet = av_packet_alloc();
		if (!vars->packet) return FALSE;
	}
	else
	{
		vars->packet = av_mallocz_av_packet();
		if (!vars->packet) return FALSE;
		av_init_packet(vars->packet);
	}

	vars->sendCount = 0;
	vars->size = frameSize;
	vars->times = Stream<Double>(8);
	vars->packets = Stream<Binary>(8);

	return TRUE;
}
void H264EncoderCPU::close()
{
	vars->close();
}
Bool H264EncoderCPU::send(CompatibleImage image, Double time)
{
	if (vars->encoder == 0 || vars->frame == 0 || vars->bgr2yuv420p == 0) return FALSE;
	if (image.width() != vars->size.width || image.height() != vars->size.height) return FALSE;
	if (image.getType() != CompatibleImageType::ImagePointerColor && image.getType() != CompatibleImageType::IplImageColor && image.getType() != CompatibleImageType::AVFrameYUV420P) return FALSE;

	Pointer target = NULL;
	if (image.getType() == CompatibleImageType::AVFrameYUV420P)
	{
		target = image.getImage();
	}
	else // CompatibleImageType::ImagePointerColor or CompatibleImageType::IplImageColor
	{
		CompatibleChannelData channelData = image.channels()[0];

		Byte *bgrSource[3] = { channelData.data(), NULL, NULL };
		Int bgrStride[3] = { (int)channelData.step(), 0, 0 };
		sws_scale(vars->bgr2yuv420p, bgrSource, bgrStride, 0, vars->size.height, av_frame_get_data(vars->frame), av_frame_get_linesize(vars->frame));

		target = vars->frame;
	}

	av_frame_set_pts(vars->frame, vars->sendCount);

	Bool ok = FALSE;
	if (getFFMPEGLibraryVersion() >= 3)
	{
		Int resultSendFrame = avcodec_send_frame(vars->encoder, target);
		ok = resultSendFrame >= 0;
	}
	else
	{
		Int gotPacket = 0;
		Int resultEncode = avcodec_encode_video2(vars->encoder, vars->packet, target, &gotPacket);
		ok = resultEncode >= 0;
		if (gotPacket)
		{
			Binary packet(av_packet_get_data(vars->packet), av_packet_get_size(vars->packet));
			vars->packets.enqueue(packet);
			av_packet_unref(vars->packet);
		}
	}

	if (ok)
	{
		vars->sendCount++;
		vars->times.enqueue(time);
	}

	return ok;
}
void H264EncoderCPU::flush()
{
	if (vars->encoder == 0) return;

	if (getFFMPEGLibraryVersion() >= 3)
	{
		avcodec_send_frame(vars->encoder, NULL);
	}
	else
	{
		Int gotPacket = 0;
		Int resultEncode = avcodec_encode_video2(vars->encoder, vars->packet, NULL, &gotPacket);
		if (resultEncode >= 0 && gotPacket)
		{
			Binary packet(av_packet_get_data(vars->packet), av_packet_get_size(vars->packet));
			vars->packets.enqueue(packet);
			av_packet_unref(vars->packet);
		}
	}
}
Bool H264EncoderCPU::receive(Binary& packetData, Double& time)
{
	if (vars->encoder == 0 || vars->packet == 0 || vars->bgr2yuv420p == 0) return FALSE;

	if (getFFMPEGLibraryVersion() >= 3)
	{
		Int resultReceive = avcodec_receive_packet(vars->encoder, vars->packet);
		if (resultReceive < 0) return FALSE;

		packetData = Binary(av_packet_get_data(vars->packet), av_packet_get_size(vars->packet));
		av_packet_unref(vars->packet);
	}
	else
	{
		Array<Binary> packets = vars->packets.dequeue(1);
		if (packets.isEmpty()) return FALSE;

		packetData = packets[0];
	}

	Array<Double> dequeued = vars->times.dequeue(1);
	Bool ok = dequeued.size() == 1;
	if (ok) time = dequeued[0];
	return ok;
}